import yaml

file_path = "/pfs/data/wangjiaqian/llm_simulator_by_alpa_layer/config/hardware_config.yaml"

def parse_hardware_config(file_path):
    """
    解析 YAML 配置文件，提取硬件设备参数信息。

    :param file_path: YAML 配置文件路径
    :return: 解析后的 Python 字典
    """
    with open(file_path, "r") as file:
        config = yaml.safe_load(file)

    # 确保配置文件包含 hardware_info
    if "hardware_info" not in config:
        raise ValueError("YAML 配置缺少 'hardware_info' 部分")

    hardware_info = config["hardware_info"]

    # 提取各部分数据
    num_nodes = hardware_info.get("num_nodes", 0)
    num_gpus_per_node = hardware_info.get("num_gpus_per_node", 0)
    total_gpus = hardware_info.get("total_gpus", num_nodes * num_gpus_per_node)

    device_mesh_info = hardware_info.get("device_mesh_info", {})
    memory = hardware_info.get("memory", {})

    # 组织解析后的数据
    intra_node_bandwidth = device_mesh_info.get('intra_node_bandwidth', 100)  # 默认100GB/s
    inter_node_bandwidth = device_mesh_info.get('inter_node_bandwidth', 12.5) # 默认12.5GB/s
    latency = device_mesh_info.get('latency', 10)  # 默认10us
    mesh_shape = device_mesh_info.get('shape', [8, 8])
    parsed_data = {
        "num_nodes": num_nodes,
        "num_gpus_per_node": num_gpus_per_node,
        "total_gpus": total_gpus,
        "mesh_shape": [num_nodes, num_gpus_per_node],
        "device_mesh": {
            "shape": mesh_shape,
            "topology": device_mesh_info.get("topology", "Unknown"),
            "intra_node_bandwidth": intra_node_bandwidth,
            "inter_node_bandwidth": inter_node_bandwidth,
            "latency": latency,
        },
        "memory": {
            "gpu_memory_per_device": memory.get("gpu_memory_per_device", None),
            "total_gpu_memory": memory.get("total_gpu_memory", None),
            "cpu_memory_per_node": memory.get("cpu_memory_per_node", None),
            "total_cpu_memory": memory.get("total_cpu_memory", None),
        }
    }
    device_mesh = (num_nodes, num_gpus_per_node)
    return parsed_data, device_mesh


# 测试函数
if __name__ == "__main__":
    hardware_config, device_mesh = parse_hardware_config(file_path)

    print("解析后的硬件信息:")
    print(device_mesh)
    print(hardware_config["num_nodes"])
